/* ****************************************************************** **
**    OpenSees - Open System for Earthquake Engineering Simulation    **
**          Pacific Earthquake Engineering Research Center            **
**                                                                    **
**                                                                    **
** (C) Copyright 1999, The Regents of the University of California    **
** All Rights Reserved.                                               **
**                                                                    **
** Commercial use of this program without express permission of the   **
** University of California, Berkeley, is strictly prohibited.  See   **
** file 'COPYRIGHT'  in main directory for information on usage and   **
** redistribution,  and for a DISCLAIMER OF ALL WARRANTIES.           **
**                                                                    **
** Developed by:                                                      **
**   Frank McKenna (fmckenna@ce.berkeley.edu)                         **
**   Gregory L. Fenves (fenves@ce.berkeley.edu)                       **
**   Filip C. Filippou (filippou@ce.berkeley.edu)                     **
**                                                                    **
** ****************************************************************** */

// Written: User-defined implementation for GPU-accelerated NewtonRaphson
// Purpose: This file contains a simplified implementation of GPU_NewtonRaphson.

#include "GPU_NewtonRaphson.h"
#include "GPU_IncrementalIntegrator.h"
#include <AnalysisModel.h>
#include <LinearSOE.h>
#include <ConvergenceTest.h>
#include <Vector.h>
#include <Channel.h>
#include <FEM_ObjectBroker.h>
#include <OPS_Globals.h>
#include <iostream>
#include <chrono>

// Constructor without test
GPU_NewtonRaphson::GPU_NewtonRaphson(int maxElements, int maxDOFs, int maxNodes)
    : NewtonRaphson(),
      gpuIntegrator(nullptr),
      gpuEnabled(true),
      profilingEnabled(false),
      maxElements(maxElements),
      maxDOFs(maxDOFs),
      maxNodes(maxNodes),
      totalIterations(0),
      totalSolveTime(0.0),
      totalFormUnbalanceTime(0.0),
      totalLinearSolveTime(0.0)
{
    initialize();
}

// Constructor with test
GPU_NewtonRaphson::GPU_NewtonRaphson(ConvergenceTest& theTest, int maxElements, int maxDOFs, int maxNodes)
    : NewtonRaphson(theTest),
      gpuIntegrator(nullptr),
      gpuEnabled(true),
      profilingEnabled(false),
      maxElements(maxElements),
      maxDOFs(maxDOFs),
      maxNodes(maxNodes),
      totalIterations(0),
      totalSolveTime(0.0),
      totalFormUnbalanceTime(0.0),
      totalLinearSolveTime(0.0)
{
    initialize();
}

// Destructor
GPU_NewtonRaphson::~GPU_NewtonRaphson()
{
    if (profilingEnabled) {
        printPerformanceStats();
    }
}

// Initialize GPU integrator
void GPU_NewtonRaphson::initialize()
{
    // Create GPU incremental integrator
    try {
        gpuIntegrator = std::make_shared<GPU_IncrementalIntegrator>(maxElements, maxDOFs, maxNodes);
        if (profilingEnabled) {
            gpuIntegrator->enableProfiling(true);
        }
    } catch (const std::exception& e) {
        opserr << "WARNING: GPU_NewtonRaphson::initialize - ";
        opserr << "failed to create GPU_IncrementalIntegrator: " << e.what() << endln;
        gpuEnabled = false;
    }
}

// Override solveCurrentStep to use GPU-accelerated integrator
int GPU_NewtonRaphson::solveCurrentStep(void)
{
    // Use parent class implementation for now - in a full implementation
    // we would intercept this to use GPU acceleration
    auto startTime = std::chrono::high_resolution_clock::now();
    int result = NewtonRaphson::solveCurrentStep();
    auto endTime = std::chrono::high_resolution_clock::now();
    
    if (profilingEnabled) {
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
        double totalTime = duration.count() / 1000.0; // Convert to milliseconds
        
        // For a real implementation, we would track actual formUnbalance and linearSolve times
        double formUnbalanceTime = 0.0;
        double linearSolveTime = 0.0;
        
        if (gpuIntegrator) {
            formUnbalanceTime = gpuIntegrator->getLastGPUTime();
        }
        
        // Estimate linear solve time as the remainder
        linearSolveTime = totalTime - formUnbalanceTime;
        if (linearSolveTime < 0.0) linearSolveTime = 0.0;
        
        // Record performance statistics
        recordPerformance(totalTime, formUnbalanceTime, linearSolveTime, 0); // We don't know iterations here
    }
    
    return result;
}

// Configuration methods
void GPU_NewtonRaphson::enableGPU(bool enable)
{
    gpuEnabled = enable;
    if (gpuIntegrator) {
        gpuIntegrator->enableGPU(enable);
    }
}

void GPU_NewtonRaphson::enableProfiling(bool enable)
{
    profilingEnabled = enable;
    if (gpuIntegrator) {
        gpuIntegrator->enableProfiling(enable);
    }
}

void GPU_NewtonRaphson::setMaxProblemSize(int elements, int dofs, int nodes)
{
    maxElements = elements;
    maxDOFs = dofs;
    maxNodes = nodes;
    
    if (gpuIntegrator) {
        gpuIntegrator->setMaxProblemSize(elements, dofs, nodes);
    }
}

// Performance monitoring
void GPU_NewtonRaphson::printPerformanceStats() const
{
    opserr << "\n------ GPU_NewtonRaphson Performance Statistics ------\n";
    opserr << "Total iterations: " << totalIterations << "\n";
    opserr << "Total solve time: " << totalSolveTime << " ms\n";
    opserr << "Total formUnbalance time: " << totalFormUnbalanceTime << " ms\n";
    opserr << "Total linear solve time: " << totalLinearSolveTime << " ms\n";
    
    if (totalIterations > 0) {
        opserr << "Average per iteration: \n";
        opserr << "  formUnbalance: " << (totalFormUnbalanceTime / totalIterations) << " ms\n";
        opserr << "  linearSolve: " << (totalLinearSolveTime / totalIterations) << " ms\n";
    }
    
    if (gpuIntegrator) {
        gpuIntegrator->printPerformanceStats();
    }
    
    opserr << "----------------------------------------------------\n";
}

// Record performance statistics
void GPU_NewtonRaphson::recordPerformance(double solveTime, double formUnbalanceTime, 
                                         double linearSolveTime, int iterations)
{
    totalSolveTime += solveTime;
    totalFormUnbalanceTime += formUnbalanceTime;
    totalLinearSolveTime += linearSolveTime;
    totalIterations += iterations;
}

// Method information
int GPU_NewtonRaphson::sendSelf(int commitTag, Channel &theChannel)
{
    // Use parent class implementation
    return NewtonRaphson::sendSelf(commitTag, theChannel);
}

int GPU_NewtonRaphson::recvSelf(int commitTag, Channel &theChannel, FEM_ObjectBroker &theBroker)
{
    // Use parent class implementation
    return NewtonRaphson::recvSelf(commitTag, theChannel, theBroker);
}

void GPU_NewtonRaphson::Print(OPS_Stream &s, int flag)
{
    // Use parent class implementation
    NewtonRaphson::Print(s, flag);
    s << "  GPU acceleration: " << (gpuEnabled ? "enabled" : "disabled") << "\n";
    s << "  Problem size limits: " << maxElements << " elements, " 
      << maxDOFs << " DOFs, " << maxNodes << " nodes\n";
}